import cv2
import os
import numpy as np

subjects = ["", "da hu zi", "Elvis Presley","hu ge"]

# 使用OpenCV用来检测脸部的函数
def detect_face(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    face_cascade = cv2.CascadeClassifier('opencv-files/lbpcascade_frontalface.xml')

    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5);

    if (len(faces) == 0):
        return None, None

    (x, y, w, h) = faces[0]

    # 只返回图像的正面部分
    return gray[y:y + w, x:x + h], faces[0]


# 该功能将读取所有人的训练图像，从每个图像检测人脸
# 并将返回两个完全相同大小的列表，一个列表
# 每张脸的脸部和另一列标签
def prepare_training_data(data_folder_path):
    # ------STEP-1--------
    # 获取数据文件夹中的目录（每个主题的一个目录）
    dirs = os.listdir(data_folder_path)

    # 列表来保存所有主题的面孔
    faces = []
    # 列表以保存所有主题的标签
    labels = []

    # 让我们浏览每个目录并阅读其中的图像
    for dir_name in dirs:

        if not dir_name.startswith("s"):
            continue;

        label = int(dir_name.replace("s", ""))

        subject_dir_path = data_folder_path + "/" + dir_name

        # 获取给定主题目录内的图像名称
        subject_images_names = os.listdir(subject_dir_path)

        # ------STEP-3--------
        # 浏览每个图片的名称，阅读图片，
        # 检测脸部并将脸部添加到脸部列表
        for image_name in subject_images_names:

            # 忽略.DS_Store之类的系统文件
            if image_name.startswith("."):
                continue;

            # 建立图像路径
            # sample image path = training-data/s1/1.pgm
            image_path = subject_dir_path + "/" + image_name

            # 阅读图像
            image = cv2.imread(image_path)

            # 显示图像窗口以显示图像
            cv2.imshow("Training on image...", image)
            cv2.waitKey(100)

            # 侦测脸部
            face, rect = detect_face(image)

            # ------STEP-4--------
            # 为了本教程的目的
            # 我们将忽略未检测到的脸部
            if face is not None:
                # 将脸添加到脸部列表
                faces.append(face)
                # 为这张脸添加标签
                labels.append(label)

    cv2.destroyAllWindows()
    cv2.waitKey(1)
    cv2.destroyAllWindows()
    print(faces,labels)
    return faces, labels

#让我们先准备好我们的训练数据
#数据将在两个相同大小的列表中
#一个列表将包含所有的面孔
#数据将在两个相同大小的列表中
print("Preparing data...")
faces, labels = prepare_training_data("training-data")
print("Data prepared")

#打印总面和标签
print("Total faces: ", len(faces))
print("Total labels: ", len(labels))

#创建我们的LBPH人脸识别器
face_recognizer = cv2.face.LBPHFaceRecognizer_create()

#或者使用EigenFaceRecognizer替换上面的行
#face_recognizer = cv2.face.createEigenFaceRecognizer()

#或者使用FisherFaceRecognizer替换上面的行
#face_recognizer = cv2.face.createFisherFaceRecognizer()

#训练我们的面部识别器
face_recognizer.train(faces, np.array(labels))


# 函数在图像上绘制矩形
# 根据给定的（x，y）坐标和
# 给定的宽度和高度
def draw_rectangle(img, rect):
    (x, y, w, h) = rect
    cv2.rectangle(img, (x, y), (x + w, y + h), (0, 255, 0), 2)


# 函数在从图像开始绘制文本
# 通过（x，y）坐标。
def draw_text(img, text, x, y):
    cv2.putText(img, text, (x, y), cv2.FONT_HERSHEY_PLAIN, 1.5, (0, 255, 0), 2)


# this function recognizes the person in image passed
# and draws a rectangle around detected face with name of the
# 学科
def predict(test_img):
    # 制作图像的副本，因为我们不想更改原始图像
    img = test_img
    # 从图像中检测脸部
    face, rect = detect_face(img)

    # 使用我们的脸部识别器预测图像
    label = face_recognizer.predict(face)
    # 获取由人脸识别器返回的相应标签的名称
    label_text = subjects[label[0]]

    # 在检测到的脸部周围画一个矩形
    draw_rectangle(img, rect)
    # 画预计人的名字
    draw_text(img, label_text, rect[0], rect[1] - 5)

    return img


print("Predicting images...")


cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("Cannot open camera")
    exit()
while True:
    # 逐帧捕获
    ret, frame = cap.read()
    # 如果正确读取帧，ret为True
    if not ret:
        print("Can't receive frame (stream end?). Exiting ...")
        break
    # 我们在框架上的操作到这里
    predicted_img1 = predict(frame)
    cv2.imshow(subjects[1], predicted_img1)
    # # 显示结果帧e
    # cv.imshow('frame', gray)
    if cv2.waitKey(1) == ord('q'):
        break
# 完成所有操作后，释放捕获器
cap.release()
cv2.destroyAllWindows()






# #加载测试图像
# # test_img1 = cv2.imread("test-data/mytest1.jpg")
# test_img1 = cv2.imread("test-data/test2.jpg")
# test_img2 = cv2.imread("test-data/test1.jpg")
# test_img3 = cv2.imread("test-data/mytest1.jpg")
#
# #执行预测
# predicted_img1 = predict(test_img1)
# predicted_img2 = predict(test_img2)
# predicted_img3 = predict(test_img3)
# print("Prediction complete")
#
# #显示两个图像
# cv2.imshow(subjects[1], predicted_img1)
# cv2.imshow(subjects[2], predicted_img2)
# cv2.imshow(subjects[3], predicted_img3)
# cv2.waitKey(0)
# cv2.destroyAllWindows()